# -*- encoding: utf-8 -*-


class Letter(object):
    def __init__(self, letter, freq):
        self.letter = letter
        self.freq = freq
        self.bitstring = ''


class TreeNode(object):
    """编码用到的小顶堆"""

    def __init__(self, freq, left, right):
        self.freq = freq
        self.left = left
        self.right = right


def parse_file(file_path):
    """处理文件内字符，计算每个字出现的次数并存入字典，最后按出现次数从小到大排序"""
    with open(file_path) as f:
        char = {}
        while True:
            c = f.read(1)
            if not c:
                break
            char[c] = char[c] + 1 if c in char else 1
    return sorted([Letter(c, f) for c, f in char.items()], key=lambda l: l.freq)


def build_tree(letters):
    """
    构建小顶堆
    """
    while len(letters) > 1:
        left = letters.pop(0)
        right = letters.pop(0)
        total_freq = left.freq + right.freq
        node = TreeNode(freq=total_freq, left=left, right=right)
        letters.append(node)
        letters.sort(key=lambda l: l.freq)
    return letters


def traverse_tree(root, bitstring):
    """
    递归编码，左子树编码+“0”，有子树编码+“1”
    递归终止条件：传入的root参数为Letter对象，即到了叶子节点
    递归公式：
        letter_list = root_left_leaf + "0" + root_right_leaf+"1"
    """
    if type(root) is Letter:
        root.bitstring = bitstring
        return [root]
    letter_list = []
    letter_list += traverse_tree(root.left, bitstring + '0')
    letter_list += traverse_tree(root.right, bitstring + '1')
    return letter_list


def huffman(file_path):
    """
    Parse the file, build the tree, then run through the file
    again, using the list of Letters to find and print out the
    bitstring for each letter.
    """
    letters_list = parse_file(file_path)
    root = build_tree(letters_list)
    letters = traverse_tree(root, "")
    print(f'Huffman Coding of {file_path}: ')
    with open(file_path) as f:
        while True:
            c = f.read(1)
            if not c:
                break
            le = list(filter(lambda l: l.letter == c, letters))[0]
            print(le.bitstring, end=" ")
    print()


if __name__ == "__main__":
    # pass the file path to the huffman function
    huffman(r'datas/file.txt')
